Skip to content

feat: [WIP/Proposal] Support copick global labels as the persisted label space#74

Open
uermel wants to merge 2 commits into
mainfrom
uermel/copick_labels
Open

feat: [WIP/Proposal] Support copick global labels as the persisted label space#74
uermel wants to merge 2 commits into
mainfrom
uermel/copick_labels

Conversation

@uermel

@uermel uermel commented Jul 1, 2026

Copy link
Copy Markdown
Collaborator

Summary

Octopi has always persisted a model-local, dense, contiguous label space
(1, 2, 3, … assigned in argument order by create-targets) into the copick
segmentations it writes (targets, predict). This PR adds an opt-in
alternative in which the segmentations octopi writes hold copick's global
object labels, while the dense [0..K] space that the model needs exists only
transiently inside the training/inference boundary.

The behavior is selected with a new --label-space flag on octopi create-targets:

  • --label-space model (default) — unchanged behavior.
  • --label-space copick — paint canonical copick global labels on disk; carry a
    {name: model_channel} map plus a label_space: copick marker through the
    targets YAML → model config → segment YAML so every stage can convert between
    the two spaces by object name.

Rationale

In copick, an integer label means the same object everywhere in a project (e.g.
gas-vesicle-shell = 2). Octopi's sequential relabeling breaks that invariant:
a targets/predict segmentation stored by octopi might relabel
gas-vesicle-shell to 1, so the file can only be interpreted if you also carry
octopi's labels map alongside it. Opening such a segmentation with the plain
project config (napari, copick logical/process ops, any downstream consumer)
mislabels every voxel.

--label-space copick fixes this at the source: what octopi writes to disk is
ordinary, canonical copick multilabel. The model still trains and predicts in a
dense [0..K] space (required by the loss, the one-hot encoding, and
RandCropByLabelClassesd), but that space is now a private, derived detail —
compacted from global labels right after loading, and expanded back to global
labels right before writing.

A second benefit: because both conversions are resolved by object name (not
by sort order or by hard-coded integers), a model trained on one project can be
applied to another project whose global labels differ, and the written
predict segmentation will use the target project's globals.

How it works

Two label concepts exist in copick mode:

  • paint label — the copick global label, written to disk.
  • model channel — the dense [0..K] index the network sees, recorded as the
    labels: {name: channel} map in the YAMLs.

The pipeline converts between them at exactly two boundaries:

  • Training (build_compaction_map + MONAI MapLabelValued, inserted right
    after LoadCopickd): global → dense. The map is derived from the recorded
    {name: channel} joined with the project's global labels by name. Objects that
    are painted but not selected as targets fold to background 0.
  • Inference (_build_dense_to_global, applied just before the single copick
    write site): dense → global, resolved by name against the target project's
    config. Metrics stay in dense space; only the write boundary expands.

localize reads the label_space marker to decide whether the predict
segmentation already holds globals (copick mode → no value remap) or the model's
dense labels (legacy → remap as before).

Examples

Assume a project with globals membrane = 1, gas-vesicle-shell = 2,
gas-vesicle = 3, and that we train on gas-vesicle-shell and gas-vesicle.

1. Default (--label-space model) — legacy, unchanged.
Targets are painted gas-vesicle-shell → 1, gas-vesicle → 2; the YAML records
labels: {gas-vesicle-shell: 1, gas-vesicle: 2} with no label_space marker.
predict is written in the same dense space. Byte-for-byte identical to before.

2. --label-space copick — the new canonical case.

  • targets seg on disk holds globals {0, 2, 3}.
  • YAML records labels: {gas-vesicle-shell: 1, gas-vesicle: 2} (dense channels,
    minted by ascending global label) and label_space: copick.
  • Training compacts {0, 2, 3} → {0, 1, 2} via orig=[0,1,2,3] → target=[0,0,1,2] (membrane is painted but unselected → 0).
  • Inference expands {0, 1, 2} → {0, 2, 3}, so the stored predict seg is
    ordinary copick multilabel and opens correctly under the project config alone.

3. Fine-tuning (--label-space copick --from-model pretrained/model_config.yaml).
Model channels are inherited per object by name from the pretrained config
instead of being minted fresh, so the channel↔object binding is preserved across
the fine-tune. If the pretrained model bound gas-vesicle → channel 5, the new
targets record channel 5 for gas-vesicle. Requesting an object the pretrained
model doesn't define is a hard error (fine-tuning onto a new object set is out of
scope).

4. Cross-project inference.
Train on project A (gas-vesicle = 3), then run octopi segment -c projectB
where gas-vesicle = 7. Because _build_dense_to_global resolves each channel
by name against project B, the written predict seg uses 7. If project B does
not define a predicted object, inference fails loudly rather than writing a
meaningless label.

Backward compatibility

The default path is untouched: without --label-space copick, no marker is
written and every code path behaves exactly as before.

For existing artifacts (targets YAMLs, model configs, and stored
segmentations produced before this change), there is no label_space marker, so:

  • get_class_info reports label_space = None;
  • build_compaction_map returns (None, None), so get_transforms skips
    MapLabelValued entirely — training is a no-op remap;
  • _build_dense_to_global falls back to identity (np.arange) — predictions are
    written in the model's dense space, unchanged;
  • localize takes the legacy global→dense remap branch.

The result is bit-identical behavior for all legacy models, targets, and
weights. No data migration is required — existing files keep working via the
identity path, and the new .get('labels') / .get('label_space') guards also
remove a latent KeyError that could occur on label-less configs.

Two consistency checks fail loudly rather than silently mislabeling:

  • Multi-config training (MultiCopickDataModule) derives a single shared
    map from the first session, so all configs must assign the same global label to
    each shared object; a disagreement raises ValueError.
  • Ensemble inference requires every model in the soup to share the same
    labels and label_space.

@daniel-ji
daniel-ji self-requested a review July 14, 2026 01:01

@daniel-ji daniel-ji left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A Claude report, reviewed and edited by me

Verdict: ✅ Sound. Core mechanism verified end-to-end; default path unchanged. Five minor, non-blocking nits.

This report can be found at [czii-dir]/cryoet-data-portal-pick-extract/octopi-74-pr/REPORT.md.

Replace [czii-dir] with the actual path to the czii dir on the HPC.


1. Environment

octopi conda env: ~/miniforge3/envs/octopi

  • pip install -e [czii-dir]/cryoet-data-portal-pick-extract/octopi --no-deps (editable branch install; regenerates the octopi console script).

Test fixture: [czii-dir]/cryoet-data-portal-pick-extract/octopi-74-pr/sample_copick_project — a copy of the 31 MB sample_copick_project unit fixture with config roots rewritten to the copy. 3 runs (TS_001–003), 64³ tomograms. Objects: proteasome=1, ribosome=2, membrane=3. Targeting ribosome alone forces divergence (global 2 ≠ dense channel 1). No original copick/octopi files were modified.


2. Test results (all pass)

Tier 0 — 19 unit tests (.../octopi-74-pr/test_tier0_units.py, no GPU): finalize_label_space (mint, divergence, --from-model inherit, missing-object error), build_compaction_map (legacy (None,None); copick orig=[0,1,2,3]/target=[0,0,1,2]), _build_dense_to_global (expansion [0,2,3], legacy identity, and the 4 raises: count mismatch, out-of-range, unknown object, ensemble disagreement). 19/19 passed.

Tier 1 — create-targets on-disk (no GPU), ribosome alone:

Mode On-disk seg (TS_001) YAML labels YAML label_space
copick [0, 2] (ribosome global) {ribosome: 1} copick
model (default) [0, 1] (dense channel) {ribosome: 1} (absent)

get_class_info + build_compaction_map on the copick target: label_space=copick, orig=[0,1,2,3] → target=[0,0,1,0] (disk 2 → channel 1). ✅

Tier 2 — GPU round-trip (A40): train → segment → localize, copick-mode ribosome target, 2 epochs, -dim 64:

  • model_config.yaml: label_space: copick, labels: {ribosome:1}, num_classes: 2. ✅
  • segment wrote predict-copick = [0, 2] on all 3 runs (global label, not channel 1) → dense→global expansion confirmed. ✅
  • localize used objects [['ribosome', 2, 150.0]] (global label kept, no legacy remap) and extracted a ribosome pick. ✅

3. Findings (minor, non-blocking)

1. --from-model on an object subset fails late, not fast.
num_classes is derived from the count of selected objects (get_class_info: len(labels)+1), but --from-model keeps each object's inherited channel. Fine-tuning on a subset can leave an inherited channel ≥ the shrunk num_classes.

Example: pretrained {ribosome→1, membrane→2} (3 classes); fine-tune on membrane onlymembrane→2 but num_classes=2 (channels {0,1}). Target value 2 is out of range; surfaces at training (one-hot out-of-bounds) or inference ("Model label 2 for 'membrane' is out of range [1, 1]"). Never silently corrupts.

Path membrane channel num_classes Result
Default mint, membrane only 1 2
--from-model, membrane only 2 2 ❌ value 2 in 2-channel model
--from-model, membrane + ribosome (empty source) 2 3 ✅ ribosome = dead channel 1

Workaround (no code change): list the full pretrained object set, giving data-less objects an empty pick source (e.g. -t "ribosome,nobody,999" -t "membrane,membrain,0"). They stay in the labels schema, so num_classes is preserved and channels stay contiguous. Verified: labels={ribosome:1, membrane:2}, num_classes=3, compaction [0,1,2,3]→[0,0,1,2], dense_to_global=[0,2,3].

Suggested fix: when --from-model is set, inherit the pretrained config's entire labels map (all objects/channels), so num_classes and channel bindings always match the pretrained model. Alternatively, add a contiguity check in finalize_label_space to fail fast.

2. uint8 ceiling. Training compaction (get_transforms(label_dtype=np.uint8)) and inference expansion (dense_to_global, np.uint8) are uint8; a global label > 255 wraps silently. Unguarded (low risk — copick segs are uint8).

3. MultiCopickDataModule missing None guard (generators.py:388): ref = {... self.roots[first_session].get_object(name).label ...} has no None check, unlike the loop at :391. First-session-missing-object → AttributeError instead of the intended ValueError. Multi-config training only.

4. get_class_info UnboundLocalError on no-seg (helpers.py:270–301): if no target segmentation is found for any runID, the return variables are never assigned. Pre-existing (loop-and-break predates the PR); warrants a clean error.

5. Compaction assumes on-disk seg holds only current pickable globals. MapLabelValued passes through any value not in orig_labels (built from root.pickable_objects). A stray on-disk label (config changed after painting) would enter dense space unmapped and could exceed num_classes. Low probability.


4. Reproduction

source ~/miniforge3/etc/profile.d/conda.sh && conda activate octopi # however you activate your octopi env
BASE=[czii-dir]/cryoet-data-portal-pick-extract/octopi-74-pr
CFG=$BASE/sample_copick_project/filesystem.json

# Tier 0
python $BASE/test_tier0_units.py

# Tier 1
octopi create-targets -c $CFG -t "ribosome,gapstop,0" -alg wbp -vs 10 \
  -name ribo-copick -uid octopi -sid 1 --label-space copick
octopi create-targets -c $CFG -t "ribosome,gapstop,0" -alg wbp -vs 10 \
  -name ribo-model  -uid octopi -sid 1

# Tier 2 (GPU)
octopi train   -c $CFG --tomo-uri wbp@10.0 --target-uri "ribo-copick:octopi/1" \
  --trainRunIDs TS_001,TS_002 --validateRunIDs TS_003 -nepochs 2 -dim 64 -bs 2 \
  -o $BASE/model_copick
octopi segment -c $CFG --tomo-uri wbp@10.0 \
  --model-config $BASE/model_copick/model_config.yaml \
  --model-weights $BASE/model_copick/best_model.pth \
  --seg-uri "predict-copick:octopi/1" --run-ids TS_001,TS_002,TS_003
octopi localize -c $CFG --seg-uri "predict-copick:octopi/1" -vs 10 \
  --runIDs TS_001,TS_002,TS_003 -pui octopi -psid loc1

Verify on-disk labels:

import copick, numpy as np
r = copick.from_file(CFG)
segs = r.get_run("TS_001").get_segmentations(name="predict-copick", user_id="octopi", session_id="1", is_multilabel=True)
print(sorted(np.unique(segs[0].numpy()).tolist()))   # -> [0, 2]  (global, not channel 1)

5. Artifacts (under [czii-dir]/cryoet-data-portal-pick-extract/octopi-74-pr/)

  • sample_copick_project/ — isolated fixture copy (roots rewritten to this dir)
  • test_tier0_units.py — 19 unit tests
  • model_copick/ — trained model (best_model.pth, model_config.yaml, results.csv)
  • train.log, seg.log, loc.log — run logs
  • pre_ribo_mb.yaml, pretrained_model_config.yaml, pretrained_missing.yaml, pre_contig.yaml--from-model fixtures

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants